# coding:utf-8
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn import tree

data_source = pd.read_excel("iris.xlsx")
data_source["属种"] = data_source["属种"].astype(np.str)
inputs = data_source[['花萼长度', '花萼宽度',
                      '花瓣长度', '花瓣宽度']].values
outputs = data_source['属种'].values
(train_inputs, test_inputs, train_outputs, test_outputs) = train_test_split(
    inputs, outputs, train_size=0.8, test_size=0.2)

# criterion="gini"
classifier = tree.DecisionTreeClassifier(criterion="entropy")
classifier.fit(train_inputs, train_outputs)
print(classifier.score(test_inputs, test_outputs))
print(classifier.predict([[0.1, 0.2, 0.3, 0.4], [0.1, 0.2, 0.3, 0.4]]))
print(classifier.predict_proba([[0.1, 0.2, 0.3, 0.4], [0.1, 0.2, 0.3, 0.4]]))
